# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_group_norm_swish")
class GroupNormSwishGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.tensor_dim = 0
        self.range_is_null = False


    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        G = random.randint(1,100)
        x_shape = case_config.inputs[0].shape
        raw_C = x_shape[1]
        C = ((raw_C + G - 1) // G) * G
        x_shape[1] = C
        case_config.inputs[0].shape = x_shape
        case_config.inputs[1].shape = [C]
        case_config.inputs[2].shape = [C]

        x_dtype = case_config.inputs[0].dtype
        case_config.inputs[1].dtype = x_dtype
        case_config.inputs[2].dtype = x_dtype
    
        case_config.inputs[3].range_values = G
        case_config.inputs[4].range_values = ["NCHW"]
        return case_config